package main

import (
	"context"
	"fmt"
	"net"
	"os"
	"shell/client/config"
	"shell/common"
	"shell/common/packages"
	"shell/global"
	"shell/utils"
	"slices"
	"strings"
	"sync"
	"time"

	"github.com/google/uuid"
	"github.com/spf13/pflag"
)

var (
	// 主要命令
	cmd = pflag.StringP("action", "a", "", "main command: shell、socks5、redir、lsport、lscend")

	// 配置文件路径
	cfgPath = pflag.StringP("config", "c", "./config.yaml", "config file path")

	// 目标被控端uuid，支持通配符
	targetUUID = pflag.StringP("uuid", "u", "", "target uuid, support wildcards")

	// shell命令
	shellCmd = pflag.StringP("shell-cmd", "s", "", "shell command")

	// 端口转发目的地址
	destAddr = pflag.StringP("dest-addr", "d", "", "port redir dest address")

	// 端口转发源端口
	srcPort = pflag.Uint16("src-port", 0, "port redir src port")

	// 帮助
	helpFlag = pflag.BoolP("help", "h", false, "show help info")

	// 端口转发动作，仅在公网转发时有效
	targetAct = pflag.StringP("target", "t", "", "port redir action, one of start soft hard")

	// 全局配置
	globalCfg *config.ClientConfig = nil

	logger = global.DefaultLogger().SetLevel(global.Debug)

	timeoutWatcher *common.TimeoutWatcher = nil
)

func main() {
	pflag.Parse()

	// 打印帮助信息
	if *helpFlag || *cmd == "" {
		fmt.Println(helpStr)
		fmt.Printf("\narguments: \n")
		pflag.PrintDefaults()
		os.Exit(0)
	}

	// 不支持的命令
	if !slices.Contains(allowCmd[:], *cmd) {
		fmt.Printf("unsupport command: %v", *cmd)
		fmt.Println(helpStr)
		fmt.Printf("\narguments: \n")
		pflag.PrintDefaults()
		os.Exit(1)
	}

	// 解析配置文件
	cfg, err := config.ParseConfig(*cfgPath)
	if err != nil {
		logger.Fatalf(0, "error parse config: %v", err)
	}
	globalCfg = cfg

	// 执行相关动作
	switch *cmd {
	case "shell":
		lqc, err := NewLazyQuicConn(globalCfg.ProxyAddr, globalCfg.TLS, 0)
		if err != nil {
			logger.Fatalf(0, "error create lazy quic conn: %v", err)
		}
		handleShell(lqc)
	case "socks5":
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		timeoutWatcher = common.NewTimeoutWatcher(ctx, 30*time.Second)
		lqc, err := NewLazyQuicConn(globalCfg.ProxyAddr, globalCfg.TLS, 60*5)
		if err != nil {
			logger.Fatalf(0, "error create lazy quic conn: %v", err)
		}
		handleSocks5(lqc)
	case "redir":
		handleRedir()
	case "lsport":
		lqc, err := NewLazyQuicConn(globalCfg.ProxyAddr, globalCfg.TLS, 0)
		if err != nil {
			logger.Fatalf(0, "error create lazy quic conn: %v", err)
		}
		handleLsPort(lqc)
	case "lscend":
		lqc, err := NewLazyQuicConn(globalCfg.ProxyAddr, globalCfg.TLS, 0)
		if err != nil {
			logger.Fatalf(0, "error create lazy quic conn: %v", err)
		}
		handleLsCend(lqc)
	}
}

// 处理列出被控端命令
func handleLsCend(lqc *LazyQuicConn) {
	stream, err := lqc.OpenStream()
	if err != nil {
		logger.Errorf(0, "error open stream: %v", err)
		return
	}
	stat, err := common.ListControlled(stream, bytePool, Role)
	if err != nil {
		common.HandleErr(err)
		stream.CloseConn(common.CloseByProcess)
		return
	}
	_ = stream.CloseConnSync(Role, true)
	fmt.Printf("------ controlled end number: %d \n", len(stat))
	for _, v := range stat {
		fmt.Printf("---- uuid: %s \n", v.UUID.String())
		fmt.Printf("     private IP: %s \n", v.PrivateIP.String())
		fmt.Printf("     public IP: %s \n", v.PublicIP.String())
		fmt.Printf("     is delay: %v \n", v.IsDelay)
		fmt.Printf("     os: %s \n", v.OperatingSystem)
		fmt.Printf("     last keep time: %s \n", v.LastKeepTime.Format(time.DateTime))
		if v.UseDelay {
			fmt.Printf("     delay uuid: %s \n", v.RelayUUID.String())
		}
	}
}

// 处理列出端口请求
func handleLsPort(lqc *LazyQuicConn) {
	req := packages.NewListGWPortReq(Role)
	stream, err := lqc.OpenStream()
	if err != nil {
		common.HandleErr(err)
		return
	}
	err = stream.WritePackage(req)
	if err != nil {
		common.HandleErr(err)
		stream.CloseConn(common.CloseByProcess)
		return
	}
	resp := new(packages.ListGWPortResp)
	err = stream.ReadPackage(resp)
	if err != nil {
		common.HandleErr(err)
		stream.CloseConn(common.CloseByProcess)
		return
	}
	serverAddr, _, _ := utils.ParseIpPort(stream.Conn.Conn.RemoteAddr().String())
	_ = stream.CloseConnSync(Role, true)
	if resp.CmdUUID != req.CmdUUID {
		logger.Fatalf(0, "error: cmd uuid not equal")
	}
	if !resp.Enabled {
		fmt.Printf("proxy end not enable public network port redir")
		return
	}
	fmt.Printf("------ port start: %d, port end: %d \n", resp.StartPort, resp.EndPort)
	fmt.Printf("------ port redir instance number: %d \n", len(resp.Stats))
	for _, v := range resp.Stats {
		fmt.Printf("---- target UUID: %s \n", v.UUID.String())
		fmt.Printf("     src addr: %s:%d \n", serverAddr.String(), v.RedirSrcPort)
		fmt.Printf("     dest address %s:%d \n", v.RedirDestIp.String(), v.RedirDestPort)
		fmt.Printf("     status: %s \n", v.Stat.String())
		fmt.Printf("     connect number: %d \n", v.ConnNum)
	}
}

// 处理socks5请求
func handleSocks5(lqc *LazyQuicConn) {
	// 判断是否启用了socks5
	if globalCfg.Socks5 == nil {
		// 配置中没有启用socks5
		logger.Fatalf(0, "socks5 not enable in config file, exit")
		return
	}
	// 启动socks5服务
	tcpListener, err := net.Listen("tcp", net.JoinHostPort("0.0.0.0", fmt.Sprintf("%d", globalCfg.Socks5.Port)))
	if err != nil {
		logger.Fatalf(0, "error start socks5 server: %v", err)
		return
	}
	defer func() {
		_ = tcpListener.Close()
	}()
	for {
		tcpConn, err := tcpListener.Accept()
		if err != nil {
			logger.Errorf(0, "error accept tcp connection: %v", err)
			break
		}
		go handleSocks5Req(tcpConn, lqc)
	}
}

// 处理单个socks5转发请求
func handleSocks5Req(conn net.Conn, lqc *LazyQuicConn) {
	req, err := common.Socks5(conn)
	if err != nil {
		logger.Errorf(0, "error get socks5 request: %v", err)
		return
	}
	var addr string
	if req.TargetAType == packages.ATypeDomain {
		addr = string(req.TargetAddr)
	} else {
		addr = net.IPv4(req.TargetAddr[0], req.TargetAddr[1], req.TargetAddr[2], req.TargetAddr[3]).String()
	}
	rule := globalCfg.Socks5.Match(addr)
	if rule == nil || rule.Action == "direct" {
		// 直接转发
		targetConn, err := net.Dial("tcp", net.JoinHostPort(addr, fmt.Sprintf("%d", req.TargetPort)))
		if err != nil {
			logger.Infof(0, "error dial tcp: %v", err)
			_ = conn.Close()
			return
		}
		common.RespSocksAndRedir(conn, targetConn, timeoutWatcher)
	} else {
		// 远程转发
		req.TargetUUID = *rule.To
		req.SetRole(Role)
		stream, err := lqc.OpenStream()
		if err != nil {
			logger.Errorf(0, "error open quic stream: %v", err)
			_ = conn.Close()
			return
		}
		// 发送socks5请求到代理端
		err = stream.WritePackage(req)
		if err != nil {
			common.HandleErr(err)
			_ = conn.Close()
			return
		}
		// 从代理端读取socks5响应
		resp := new(packages.Socks5RedirResp)
		err = stream.ReadPackage(resp)
		if err != nil {
			common.HandleErr(err)
			_ = conn.Close()
			return
		}
		if !resp.IsOk {
			logger.Infof(0, "error socks5 resp: %s", string(resp.ErrorInfo))
			_ = conn.Close()
			_ = stream.CloseStreamSync(Role, true)
			return
		}
		if resp.CmdUUID != req.CmdUUID {
			logger.Errorf(0, "error socks5 resp uuid")
			_ = conn.Close()
			_ = stream.CloseStreamSync(Role, true)
			return
		}
		common.RespSocksAndRedir(conn, stream, timeoutWatcher)
	}
}

// 处理端口转发请求
func handleRedir() {
	// 处理参数，确定是本地转发还是公网转发
	if *targetAct == "" {
		// 是本地转发
		// 如果指定了公网转发相关参数，就退出
		if *targetUUID != "" || *srcPort != 0 || *destAddr != "" {
			logger.Fatalf(0, "local port redir not need args: --uuid --src-port --dest-addr")
		}
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		timeoutWatcher = common.NewTimeoutWatcher(ctx, 30*time.Second)
		lqc, err := NewLazyQuicConn(globalCfg.ProxyAddr, globalCfg.TLS, -1)
		if err != nil {
			logger.Fatalf(0, "error create lazy quic conn: %v", err)
		}
		handleLocalPortRedir(lqc)
	} else {
		// 是公网转发
		if !slices.Contains(allowTarget[:], *targetAct) {
			// 目标动作不合法
			logger.Errorf(0, "error target action of public network port redir")
			logger.Errorf(0, "      allow action is : %v", allowTarget)
		}
		// 解析参数
		if *targetUUID == "" || *srcPort == 0 {
			logger.Fatalf(0, "set target uuid and source port for start local port redir")
		}
		destIP, destPort, err := utils.ParseIpPort(*destAddr)
		if err != nil {
			logger.Fatalf(0, "dest addr format error: %v", err)
		}
		lqc, err := NewLazyQuicConn(globalCfg.ProxyAddr, globalCfg.TLS, 0)
		if err != nil {
			logger.Fatalf(0, "error create lazy quic conn: %v", err)
		}
		stream, err := lqc.OpenStream()
		if err != nil {
			logger.Fatalf(0, "error open quic stream: %v", err)
		}
		stat, err := common.ListControlled(stream, bytePool, Role)
		if err != nil {
			logger.Errorf(0, "error list controlled end: %v", err)
			stream.CloseConn(common.CloseByProcess)
			return
		}
		matchStat := common.MatchUUID(stat, *targetUUID)
		switch len(matchStat) {
		case 0:
			// 没有匹配的被控端
			logger.Errorf(0, "not match any controlled end, exit")
			_ = stream.CloseConnSync(Role, true)
			return
		case 1:
			req := packages.NewPortRedirReq(true, PortAction(*targetAct), matchStat[0].UUID, destIP, destPort, *srcPort, Role)
			handleGWPortRedir(req, stream)
		default:
			// 匹配到多个被控端
			logger.Errorf(0, "match more than one controlled end, exit")
			_ = stream.CloseConnSync(Role, true)
			for _, s := range matchStat {
				fmt.Printf("%s \n", s.UUID.String())
			}
			return
		}
	}
}

// 处理公网端口转发，完成后关闭quic连接
func handleGWPortRedir(rep *packages.PortRedirReq, stream *common.QuicStream) {
	err := stream.WritePackage(rep)
	if err != nil {
		common.HandleErr(err)
		stream.CloseConn(common.CloseByProcess)
		return
	}
	resp := new(packages.PortRedirResp)
	err = stream.ReadPackage(resp)
	if err != nil {
		common.HandleErr(err)
		stream.CloseConn(common.CloseByProcess)
		return
	}
	if resp.IsOk {
		logger.Infof(0, "ok")
	} else {
		logger.Errorf(0, "error: %s", string(resp.ErrorInfo))
	}
	_ = stream.CloseConnSync(Role, true)
}

// 处理本地端口转发
func handleLocalPortRedir(lqc *LazyQuicConn) {
	// 解析配置文件
	l := len(globalCfg.Redir)
	switch l {
	case 0:
		logger.Fatalf(0, "error no redir config in config file")
	case 1:
		redir := &(globalCfg.Redir[0])
		destIP, destPort, _ := utils.ParseIpPort(redir.RedirAddr)
		req := packages.NewPortRedirReq(false, packages.PortOpen, redir.UUID, destIP, destPort, redir.LocalPort, Role)
		redirInst := NewPortRedirInst(req, lqc, bytePool)
		if err := redirInst.StartRedir(); err != nil {
			logger.Errorf(0, "error start local port redir instance: %v", err)
		}
		redirInst.Stop()
	default:
		wg := sync.WaitGroup{}
		for _, v := range globalCfg.Redir {
			destIP, destPort, _ := utils.ParseIpPort(v.RedirAddr)
			req := packages.NewPortRedirReq(false, packages.PortOpen, v.UUID, destIP, destPort, v.LocalPort, Role)
			redirInst := NewPortRedirInst(req, lqc, bytePool)
			wg.Add(1)
			go func() {
				if err := redirInst.StartRedir(); err != nil {
					logger.Errorf(0, "error start local port redir instance: %v", err)
				}
				redirInst.Stop()
				wg.Done()
			}()
		}
		wg.Wait()
	}
}

// 处理执行shell命令请求
func handleShell(lqc *LazyQuicConn) {
	if *targetUUID == "" {
		logger.Fatalf(0, "error unset target uuid")
	}
	if *shellCmd == "" {
		logger.Fatalf(0, "error unset shell command")
	}
	conReq := packages.NewControlledReq(Role)
	stream, err := lqc.OpenStream()
	if err != nil {
		common.HandleErr(err)
		return
	}
	err = stream.WritePackage(conReq)
	if err != nil {
		common.HandleErr(err)
		stream.CloseConn(common.CloseByProcess)
		return
	}
	conResp := new(packages.ControlledResp)
	err = stream.ReadPackage(conResp)
	if err != nil {
		common.HandleErr(err)
		stream.CloseConn(common.CloseByProcess)
		return
	}
	if len(conResp.Stats) == 0 {
		logger.Infof(0, "controlled end is down")
		_ = stream.CloseConnSync(Role, true)
		return
	}
	var fullUUID uuid.UUID
	var matchNum int = 0
	for _, v := range conResp.Stats {
		if strings.HasPrefix(v.UUID.String(), *targetUUID) {
			fullUUID = v.UUID
			matchNum++
		}
	}
	if matchNum == 0 {
		logger.Errorf(0, "input uuid not match any controlled end")
		_ = stream.CloseConnSync(Role, true)
		return
	}
	if matchNum > 1 {
		logger.Errorf(0, "input uuid match more than one controlled end")
		_ = stream.CloseConnSync(Role, true)
		return
	}
	shellReq := new(packages.ShellCmdReq)
	shellReq.CmdUUID = uuid.New()
	shellReq.TargetUUID = fullUUID
	shellReq.ShellCmd = []byte(*shellCmd)
	shellReq.SetRole(Role)
	err = stream.WritePackage(shellReq)
	if err != nil {
		common.HandleErr(err)
		stream.CloseConn(common.CloseByProcess)
		return
	}
	shellResp := new(packages.ShellCmdResp)
	err = stream.ReadPackage(shellResp)
	if err != nil {
		common.HandleErr(err)
		stream.CloseConn(common.CloseByProcess)
		return
	}
	fmt.Printf("command exit code: %d \n", shellResp.CmdExitCode)
	fmt.Printf("command std out: \n%s \n", string(shellResp.CmdOutput))
	_ = stream.CloseConnSync(Role, true)
}
